/**
 * Copyright (c) 2022 Presto Labs Pte. Ltd.
 * Author: Ju
 */
// Base class for packing & unpacking protocols xmitted through networks

#pragma once

#include <algorithm>
#include <cassert>
#include <cstdint>
#include <cstring>
#include <functional>
#include <list>
#include <map>
#include <memory>
#include <set>
#include <stdexcept>
#include <string>
#include <unordered_map>
#include <unordered_set>
#include <utility>
#include <vector>

#include "coin2/base/cpp_feature.h"
#include "coin2/base/memory/mem_buf.h"
#include "coin2/base/memory/mem_slice.h"

namespace coin2::base {
using namespace coin2::base::memory;

enum class PacketLimit : uint32_t {
  kMaxPacketSize = 8 * 1024 * 1024 - 1024,
};

class Packer : private base::NonCopyOrMovable {
 public:
  enum {
    kDefaultSize = 1024,
  };

  explicit Packer(uint32_t buffer_sz = kDefaultSize);

  // After |Pack| is called. All |Pushxxx| functions should not be called
  // until |Reset| is called.
  Packer &Pack();
  void Reset();

  Packer &Pack2();

  void PushBoolean(bool val);
  void PushUint8(uint8_t val);
  void PushUint16(uint16_t val);
  void PushUint32(uint32_t val);
  void PushUint64(uint64_t val);
  void PushFloat(float val);
  void PushDouble(double val);

  void PushString(const std::string &val);
  void PushBuffer(const std::vector<uint8_t> &val);
  void PushBinary(const void *data, size_t length);

  const char *Buffer() const;
  uint32_t Length() const;

  // The following two functions are valid only after |Pack| is called.
  std::string Data() const;

  memory::MemSlice TakeBuffer();

 private:
  struct MemBufDeleter {
    void operator()(memory::MemBuf *buf) const { buf->Release(); }
  };

  typedef std::unique_ptr<memory::MemBuf, MemBufDeleter> MemBufPtr;

  void CheckSize(uint32_t more, uint32_t position);

  char *Begin();
  const char *Begin() const;
  char *End();

 private:
  MemBufPtr mem_;
  uint32_t length_;
  uint32_t position_;
};

inline Packer::Packer(uint32_t buffer_sz) {
  mem_.reset(memory::MemBuf::Create(std::max<uint32_t>(64u, buffer_sz)));

  length_ = 0;
  position_ = 4;
}

inline Packer &Packer::Pack() {
  assert(position_ < static_cast<uint32_t>(PacketLimit::kMaxPacketSize));

  length_ = position_;
  position_ = 0;

  PushUint32(length_);
  position_ = 0;
  return *this;
}

// The header length is not added
inline Packer &Packer::Pack2() {
  assert(position_ < static_cast<uint32_t>(PacketLimit::kMaxPacketSize));

  length_ = position_ - 4;
  position_ = 4;
  return *this;
}

inline void Packer::Reset() {
  length_ = 0;
  position_ = 4;
  if (!mem_) {
    mem_.reset(memory::MemBuf::Create(64));
  }
}

inline void Packer::PushBoolean(bool val) {
  char cval = static_cast<char>(val);
  PushUint8(cval);
}

inline void Packer::PushUint8(uint8_t val) {
  CheckSize(sizeof(val), position_);

  ::memcpy(Begin() + position_, &val, sizeof(val));
  position_ += sizeof(val);
}

inline void Packer::PushUint16(uint16_t val) {
  CheckSize(sizeof(val), position_);

  ::memcpy(Begin() + position_, &val, sizeof(val));
  position_ += sizeof(val);
}

inline void Packer::PushUint32(uint32_t val) {
  CheckSize(sizeof(val), position_);

  ::memcpy(Begin() + position_, &val, sizeof(val));
  position_ += sizeof(val);
}

inline void Packer::PushUint64(uint64_t val) {
  CheckSize(sizeof(val), position_);

  ::memcpy(End(), &val, sizeof(val));
  position_ += sizeof(val);
}

inline void Packer::PushFloat(float val) {
  CheckSize(sizeof(val), position_);
  ::memcpy(End(), &val, sizeof(val));
  position_ += sizeof(val);
}

inline void Packer::PushDouble(double val) {
  CheckSize(sizeof(val), position_);

  ::memcpy(End(), &val, sizeof(val));
  position_ += sizeof(val);
}

inline void Packer::PushString(const std::string &val) {
  uint32_t length = static_cast<uint32_t>(val.length());
  PushUint32(length);

  CheckSize(length, position_);

  if (length > 0) {
    ::memcpy(End(), val.data(), length);
    position_ += length;
  }
}

inline void Packer::PushBuffer(const std::vector<uint8_t> &val) {
  uint32_t length = static_cast<uint32_t>(val.size());
  PushUint32(length);

  CheckSize(length, position_);

  if (length > 0) {
    ::memcpy(End(), val.data(), length);
    position_ += length;
  }
}

inline void Packer::PushBinary(const void *data, size_t length) {
  CheckSize(length, position_);

  if (length > 0) {
    ::memcpy(End(), data, length);
    position_ += static_cast<uint32_t>(length);
  }
}

inline const char *Packer::Buffer() const { return Begin() + position_; }

inline uint32_t Packer::Length() const { return length_; }

inline std::string Packer::Data() const {
  return std::string(Begin() + position_, length_);
}

inline void Packer::CheckSize(uint32_t more, uint32_t position) {
  uint32_t new_size = position + more;

  if (new_size >= static_cast<uint32_t>(PacketLimit::kMaxPacketSize)) {
    throw std::overflow_error("packer buffer overflow!");
  }

  auto old_cap = mem_->Capacity();
  if (old_cap < new_size) {
    auto size = std::max<uint32_t>(new_size, 2 * old_cap);
    mem_.reset(memory::MemBuf::Create(mem_->Begin(), old_cap, size));
  }
}

inline const char *Packer::Begin() const {
  return reinterpret_cast<const char *>(mem_->Begin());
}

inline char *Packer::Begin() { return reinterpret_cast<char *>(mem_->Begin()); }

inline char *Packer::End() { return Begin() + position_; }

inline memory::MemSlice Packer::TakeBuffer() {
  memory::MemSlice slice(mem_.release());
  slice.SetBeginPointer(position_);
  slice.SetEndPointer(position_ + length_);
  return slice;
}

inline Packer &operator<<(Packer &p, bool v) {
  p.PushBoolean(v);
  return p;
}

inline Packer &operator<<(Packer &p, uint8_t v) {
  p.PushUint8(v);
  return p;
}

inline Packer &operator<<(Packer &p, uint16_t v) {
  p.PushUint16(v);
  return p;
}

inline Packer &operator<<(Packer &p, uint32_t v) {
  p.PushUint32(v);
  return p;
}

inline Packer &operator<<(Packer &p, uint64_t v) {
  p.PushUint64(v);
  return p;
}

inline Packer &operator<<(Packer &p, char v) {
  p.PushUint8(static_cast<uint8_t>(v));
  return p;
}

inline Packer &operator<<(Packer &p, int8_t v) {
  p.PushUint8(static_cast<uint8_t>(v));
  return p;
}

inline Packer &operator<<(Packer &p, int16_t v) {
  p.PushUint16(static_cast<uint16_t>(v));
  return p;
}

inline Packer &operator<<(Packer &p, int32_t v) {
  p.PushUint32(static_cast<uint32_t>(v));
  return p;
}

inline Packer &operator<<(Packer &p, int64_t v) {
  p.PushUint64(static_cast<uint64_t>(v));
  return p;
}

inline Packer &operator<<(Packer &p, float v) {
  p.PushFloat(static_cast<float>(v));
  return p;
}

inline Packer &operator<<(Packer &p, double v) {
  p.PushDouble(static_cast<double>(v));
  return p;
}

inline Packer &operator<<(Packer &p, const std::vector<uint8_t> &v) {
  p.PushBuffer(v);
  return p;
}

inline Packer &operator<<(Packer &p, const std::string &v) {
  p.PushString(v);
  return p;
}

template <typename T>
inline Packer &operator<<(Packer &p, const std::vector<T> &v) {
  uint32_t count = static_cast<uint32_t>(v.size());
  p.PushUint32(count);

  for (uint32_t i = 0; i < count; ++i) {
    p << v[i];
  }

  return p;
}

template <typename T>
inline Packer &operator<<(Packer &p, const std::list<T> &v) {
  uint32_t count = static_cast<uint32_t>(v.size());
  p.PushUint32(count);

  for (const T &x : v) {
    p << x;
  }

  return p;
}

template <typename T>
inline Packer &operator<<(Packer &p, const std::set<T> &v) {
  uint32_t count = static_cast<uint32_t>(v.size());
  p.PushUint32(count);

  for (const T &x : v) {
    p << x;
  }

  return p;
}

template <typename K, typename V>
inline Packer &operator<<(Packer &p, const std::pair<K, V> &value) {
  p << value.first << value.second;
  return p;
}

template <typename K, typename V>
inline Packer &operator<<(Packer &p, const std::map<K, V> &v) {
  uint32_t count = static_cast<uint32_t>(v.size());
  p.PushUint32(count);

  for (const auto &x : v) {
    p << x;
  }

  return p;
}

template <typename K, typename V>
inline Packer &operator<<(Packer &p, const std::unordered_map<K, V> &v) {
  uint32_t count = static_cast<uint32_t>(v.size());
  p.PushUint32(count);

  for (const auto &x : v) {
    p << x;
  }

  return p;
}

template <typename T, size_t Len>
inline Packer &operator<<(Packer &p, const std::array<T, Len> &arr) {
  for (const auto &x : arr) {
    p << x;
  }
  return p;
}

template <typename K>
inline Packer &operator<<(Packer &p, const std::unordered_set<K> &s) {
  uint32_t count = static_cast<uint32_t>(s.size());
  p.PushUint32(count);

  for (const auto &x : s) {
    p << x;
  }

  return p;
}

}  // namespace coin2::base
